import os
import time

import torch
import torchvision
from torch import nn
from torch.autograd import Variable
from torch.utils.data import DataLoader
from torchvision import transforms
from torchvision.datasets import MNIST, CIFAR10
from torchvision.utils import save_image
import numpy as np
import pandas as pd
import torch
import os
from datetime import datetime
import matplotlib.pyplot as plt
from torch.utils.data import DataLoader
from torch.autograd import Variable
from torch import nn
from torch import optim
from torchvision import datasets



class AutoEncoder(nn.Module):

    def __init__(self, inputSize, hiddenSize):
        super().__init__()
        self.inputSize = inputSize
        self.hiddenSize = hiddenSize

        self.encoder = nn.Linear(inputSize, hiddenSize, bias=True)
        self.decoder = nn.Linear(hiddenSize, inputSize, bias=True)

        #self.activation = torch.nn.LeakyReLU()
        self.activation = torch.nn.ReLU()

    def forward(self, x, isHidden=False):

        hidden = self.activation(self.encoder(x))
        if isHidden:
            return hidden
        else:
            return self.activation(self.decoder(hidden))



class SAE(nn.Module):
    def __init__(self, encoderList):
        super().__init__()
        self.ae1 = encoderList[0]
        self.ae2 = encoderList[1]


    def forward(self, x):
        a1 = self.ae1(x, isHidden=True)
        a2 = self.ae2(a1, isHidden=True)
        out = self.pre1(a2)
        return out




